#pragma once

#if USE_VULKAN
#include <vulkan/vulkan.h>
#include <memory>
#include "vkinfer/runtime/context.h"
#include "vktransfer.h"
#include "vkcompute.h"

namespace vkinfer
{
	class VkOp;
	class VkContext : public Context
	{
	public:
		VkInstance instance = VK_NULL_HANDLE;
		// VkDebugReportCallbackEXT debugReportCallback = VK_NULL_HANDLE;
		VkPhysicalDevice gpu = VK_NULL_HANDLE;
		VkDevice device = VK_NULL_HANDLE;
		
		VkQueue compute_queue = VK_NULL_HANDLE;
		uint32_t compute_queue_index = 0;
		VkQueue transfer_queue = VK_NULL_HANDLE;
		uint32_t transfer_queue_index = 0;

		VkFence fence;

		// multiple instances may allow multi-threaded computation.
		std::shared_ptr<VkCompute> compute;
		std::shared_ptr<VkTransfer> transfer;

	public:
		VkContext();
		~VkContext();

		void begin() override;
		void end() override;
		void submit() override;

		std::shared_ptr<Blob> make_blob(const std::string& name, BlobTransferDirection direction) override;
		std::shared_ptr<Blob> make_blob(const std::string& name, uint32_t count, BlobTransferDirection direction) override;
		std::shared_ptr<Blob> make_blob(const std::string& name, const std::vector<uint32_t>& shape, BlobTransferDirection direction) override;
		std::shared_ptr<Operator> make_op(const std::string& name, const std::shared_ptr<OnnxNodeInfo>& node_info) override;
	
	private:
		// initialization functions
		void create_instance();
		void pick_physical_device();
		void create_device();
	};
}


#endif
